# -*- coding: utf-8 -*-
import random
import time
from typing import List

import cv2
import gym
import gym.spaces as spaces
# 导入按键的检测
from pygame.locals import *
import numpy as np

from ril_env.envs.space_war.space_war_util import *


class SpaceWarEnv(gym.core.Env, Base):
	enemy_frequency = 0  # 敌机产生频率
	DO_NOTHING = Player.DO_NOTHING
	MOVE_UP = Player.MOVE_UP
	MOVE_DOWN = Player.MOVE_DOWN
	MOVE_LEFT = Player.MOVE_LEFT
	MOVE_RIGHT = Player.MOVE_RIGHT
	SHOOT = Player.SHOOT
	SCREEN_WIDTH, SCREEN_HEIGHT = SCREEN_WIDTH, SCREEN_HEIGHT

	def __init__(self):
		pygame.display.set_caption("Space War")
		self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
		pygame.init()
		super().__init__('background.png')

		self.player = Player()
		self.enemyList: List[Enemy] = []
		self.observation_space = spaces.Box(
			low=0, high=255, shape=(4, SCREEN_WIDTH, SCREEN_HEIGHT), dtype=np.int
		)
		self.action_space = spaces.Discrete(6)

	def step(self, action: int):
		pygame.event.pump()  # 内部处理pygame事件处理程序
		terminate, reward = False, 0.1
		if not self.player.alive:  # 滞后一回合结束，为了让飞机 crash 的图像能够被看到
			terminate = True
			reward = -5  # 原游戏是 -1
		else:
			self.handle_action(action)
			self.gen_enemy()

			for enemy in self.enemyList:
				if enemy.check_out() or not enemy.alive:  # 清除敌人
					self.enemyList.remove(enemy)
					continue
				enemy.move()

				for bullet in self.player.bullets:
					if enemy.x <= (bullet.x + bullet.WIDTH // 2) <= (enemy.x + enemy.WIDTH) \
							and bullet.y <= (enemy.y + enemy.HEIGHT):
						enemy.crash()
						bullet.alive = False
						reward += 1  # 消灭一个敌人得一分（原游戏是无论消灭几个都只是1）
						break

				if self.player.x <= (enemy.x + enemy.WIDTH // 2) <= (self.player.x + self.player.WIDTH) \
						and self.player.y <= (enemy.y + enemy.HEIGHT):
					self.player.crash()
					enemy.crash()

			self.player.clear_bullet()
			for bullet in self.player.bullets:
				bullet.move()
			self.update_frame()

		observation = pygame.surfarray.array3d(pygame.display.get_surface())
		return observation, reward, terminate, {}

	def gen_enemy(self):
		"""产生敌方飞机"""
		if self.enemy_frequency % 30 == 0:
			rand_x = random.randint(0, SCREEN_WIDTH - 51)
			enemy = Enemy(rand_x)
			self.enemyList.append(enemy)
		self.enemy_frequency += 1
		if self.enemy_frequency >= 100:
			self.enemy_frequency = 0

	def update_frame(self):
		self.draw(self.screen)
		self.player.draw(self.screen)
		for bullet in self.player.bullets:
			bullet.draw(self.screen)
		for enemy in self.enemyList:
			enemy.draw(self.screen)

	def reset(self):
		self.__init__()
		return self.step(self.DO_NOTHING)[0]

	def seed(self, seed=None):
		return []

	def render(self, mode='human'):
		pygame.display.update()

	def close(self):
		pygame.quit()

	def handle_action(self, action):
		if action in [self.MOVE_UP, self.MOVE_DOWN, self.MOVE_LEFT, self.MOVE_RIGHT]:
			self.player.move(action)
		elif action == self.SHOOT:
			self.player.shoot()
		else:
			pass

	@staticmethod
	def get_action_meanings():
		return ['NOOP', 'MOVE_UP', 'MOVE_DOWN', 'MOVE_LEFT', 'MOVE_RIGHT', 'SHOOT']

	def human_play(self, show_obs=False):
		terminate = False
		self.reset()
		score = 0
		while not terminate:
			act = self.DO_NOTHING
			# 判断按键
			for event in pygame.event.get():
				if event.type == QUIT:
					exit()
				elif event.type == KEYDOWN:
					if event.key == K_w or event.key == K_UP:
						act = self.MOVE_UP
					elif event.key == K_s or event.key == K_DOWN:
						act = self.MOVE_DOWN
					elif event.key == K_a or event.key == K_LEFT:
						act = self.MOVE_LEFT
					elif event.key == K_d or event.key == K_RIGHT:
						act = self.MOVE_RIGHT
					elif event.key == K_SPACE:
						act = self.SHOOT
			observation, reward, terminate, _ = self.step(act)
			score += reward
			self.render()
			time.sleep(0.03)
		print('score: %f' % score)
